package com.rmfyzxfw.caseprocess.config.datasource;

import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;

import javax.sql.DataSource;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

public class DataSourceRouting extends AbstractRoutingDataSource {

    private final Map<Object, DataSource> outerDataSources = new ConcurrentHashMap<>();

    @Override
    protected Object determineCurrentLookupKey() {
        return DataSourceContextHolder.getDataSourceKey();
    }

    @Override
    public DataSource determineTargetDataSource() {
        Object lookupKey = determineCurrentLookupKey();
        if (lookupKey == null) {
            return super.determineTargetDataSource();
        }

        DataSource dataSource = this.outerDataSources.get(lookupKey);
        if (dataSource != null) {
            return dataSource;
        }
        return super.determineTargetDataSource();
    }

    public void addOuterDataSource(Object key, DataSource dataSource) {
        this.outerDataSources.put(key, dataSource);
    }

    public void removeOuterDataSource(Object key) {
        DataSource dataSource = outerDataSources.get(key);
        if (dataSource == null) {
            return;
        }
        outerDataSources.remove(key);
    }

    public boolean outerDataSourceContainsKey(Object key) {
        return outerDataSources.containsKey(key);
    }
}
